"""
MongoDB数据库操作模块
"""
from pymongo import MongoClient
from pymongo.errors import ConnectionFailure, OperationFailure
from config import MONGODB_CONFIG
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class Database:
    def __init__(self):
        self.client = None
        self.db = None
        self.connect()

    def connect(self):
        """连接到MongoDB"""
        try:
            connection_string = f"mongodb://{MONGODB_CONFIG['username']}:{MONGODB_CONFIG['password']}@{MONGODB_CONFIG['host']}:{MONGODB_CONFIG['port']}/{MONGODB_CONFIG['database']}"
            self.client = MongoClient(connection_string, serverSelectionTimeoutMS=5000)
            self.db = self.client[MONGODB_CONFIG['database']]
            # 测试连接
            self.client.admin.command('ping')
            logger.info("MongoDB连接成功")
        except (ConnectionFailure, OperationFailure) as e:
            logger.error(f"MongoDB连接失败: {e}")
            raise

    def get_collection(self, collection_name):
        """获取集合"""
        return self.db[collection_name]

    def close(self):
        """关闭连接"""
        if self.client:
            self.client.close()
            logger.info("MongoDB连接已关闭")


# 账号管理
class AccountManager:
    def __init__(self, db: Database):
        self.collection = db.get_collection('accounts')

    def create_account(self, username, password):
        """创建账号"""
        if self.collection.find_one({'username': username}):
            return False, "用户名已存在"
        account = {
            'username': username,
            'password': password,  # 实际应用中应该加密存储
            'status': 'active',
            'created_at': None
        }
        self.collection.insert_one(account)
        return True, "账号创建成功"

    def verify_account(self, username, password):
        """验证账号"""
        account = self.collection.find_one({'username': username, 'password': password})
        return account is not None

    def get_account(self, username):
        """获取账号信息"""
        return self.collection.find_one({'username': username})

    def update_account_status(self, username, status):
        """更新账号状态"""
        self.collection.update_one(
            {'username': username},
            {'$set': {'status': status}}
        )

    def list_accounts(self):
        """列出所有账号"""
        return list(self.collection.find({}, {'password': 0}))


# 客户端管理
class ClientManager:
    def __init__(self, db: Database):
        self.collection = db.get_collection('clients')

    def add_client(self, client_id, username, ip_address):
        """添加客户端"""
        client = {
            'client_id': client_id,
            'username': username,
            'ip_address': ip_address,
            'status': 'online',
            'whitelist': False,
            'blacklist': False,
            'last_seen': None
        }
        self.collection.insert_one(client)
        return True

    def update_client_status(self, client_id, status):
        """更新客户端状态"""
        self.collection.update_one(
            {'client_id': client_id},
            {'$set': {'status': status, 'last_seen': None}}
        )

    def set_whitelist(self, client_id, whitelist):
        """设置白名单"""
        self.collection.update_one(
            {'client_id': client_id},
            {'$set': {'whitelist': whitelist}}
        )

    def set_blacklist(self, client_id, blacklist):
        """设置黑名单"""
        self.collection.update_one(
            {'client_id': client_id},
            {'$set': {'blacklist': blacklist}}
        )

    def get_client(self, client_id):
        """获取客户端信息"""
        return self.collection.find_one({'client_id': client_id})

    def list_clients(self):
        """列出所有客户端"""
        return list(self.collection.find({}))

    def get_online_clients(self):
        """获取在线客户端"""
        return list(self.collection.find({'status': 'online'}))

    def remove_client(self, client_id):
        """删除客户端"""
        self.collection.delete_one({'client_id': client_id})

